import numba
import numpy as np
import config
import image_feature
from matplotlib import pyplot as plt
import cv2
import tools


class ImageVectorExtractor():
    def __init__(self, path1, path2):
        self.path1 = path1
        self.path2 = path2
        self.img1 = cv2.imread(path1, cv2.COLOR_BGR2GRAY)
        self.img2 = cv2.imread(path2, cv2.COLOR_BGR2GRAY)

    # @numba.jit
    def match(self):
        self.k1, self.d1 = image_feature.detect(self.img1)
        self.k2, self.d2 = image_feature.detect(self.img2)

        matcher = cv2.BFMatcher()
        matches = matcher.knnMatch(self.d1, self.d2, k=2)

        '''查找最佳匹配'''
        goodMatches, matchesMask = image_feature.bf_knnmatches(matches, self.k1, self.k2)
        self.matched = []
        for match, mask in zip(goodMatches, matchesMask):
            if mask[0] == 1:
                self.matched.append(match)

        '''找出两副图片的匹配点'''
        self.matched_points = []
        for m in self.matched:
            self.matched_points.append((self.k1[m[0].queryIdx], self.k2[m[0].trainIdx]))

    # @numba.jit
    def extract_vectors(self):
        self.match()
        self.speed_matrix = np.zeros((self.img1.shape[0], self.img1.shape[1], 2), np.float32)
        self.calc_matched_vectors()

        # start = time.time()
        # self.calc_not_match_vectors()
        # end = time.time()
        # print(end - start)
        #
        # print(self.speed_matrix.shape)
        # np.save('vector_matrix', self.speed_matrix)

    # @numba.jit
    def calc_matched_vectors(self):
        '''计算已经匹配的特征点的运功矢量'''

        self.matchedPointVectors = []

        for (p1, p2) in self.matched_points:
            speed_x = p2.pt[0] - p1.pt[0]
            speed_y = p2.pt[1] - p1.pt[1]

            x = int(p2.pt[0])
            y = int(p2.pt[1])
            vector = self.speed_matrix[y][x]

            vector[0] = speed_x
            vector[1] = speed_y
            # vector[2] = self.img2[y][x]

            # print(vector)

            self.matchedPointVectors.append((x, y, vector))

    # @numba.jit
    def findNearerMatchedPoint(self, point, num=3):

        points = []
        for pv in self.matchedPointVectors:
            distance = tools.point_distance(point, pv)
            points.append({'distance': distance, 'pv': pv})

        points.sort(key=lambda x: x['distance'])
        nearer = []
        for x in points[:num]:
            nearer.append(x['pv'])

        return nearer

    # @numba.jit
    def nearer_vector(self, nearer):
        '''邻近点的速度计算平均速度'''

        speed_x_total = 0
        speed_y_total = 0
        for p in nearer:
            # 从临近点速度计算出平均速度
            speed_x = p[2][0]
            speed_y = p[2][1]

            speed_x_total += speed_x
            speed_y_total += speed_y

        speed_x_avg = speed_x_total / len(nearer)
        speed_y_avg = speed_y_total / len(nearer)

        return speed_x_avg, speed_y_avg

    # @numba.jit
    def calc_not_match_vectors(self):
        '''计算未匹配的点的运功矢量'''

        print('计算未匹配的点的运功矢量')
        for y in range(self.speed_matrix.shape[0]):
            print(y)
            for x in range(self.speed_matrix.shape[1]):
                vector = self.speed_matrix[x][y]
                radarValue = self.img2[x][y]
                if vector[0] == 0:  # 未匹配
                    nearer = self.findNearerMatchedPoint((x, y))
                    avg = self.nearer_vector(nearer)
                    vector[0] = avg[0]
                    vector[1] = avg[1]
                    vector[2] = radarValue

    def showImg1MatchedPoints(self):
        points = []
        for p in self.matched_points:
            points.append(p[0])

        image = cv2.drawKeypoints(self.img1, points, self.img1)
        cv2.imshow('matchedPoints', image)
        cv2.waitKey()

    def showImg2MatchedPoints(self):
        points = []
        for p in self.matched_points:
            points.append(p[1])

        image = cv2.drawKeypoints(self.img2, points, self.img2)
        cv2.imshow('matchedPoints', image)
        cv2.waitKey()


def show_hot():
    speed_matrix_list = []

    speed_x = []

    for i in range(1):
        imageVectorExactor = ImageVectorExtractor(config.BLOCK + '%02d.png' % (i,),
                                                  config.BLOCK + '%02d.png' % (i + 1,))

        print("process :", imageVectorExactor.path1 + "," + imageVectorExactor.path2)

        imageVectorExactor.match()
        # imageVectorExactor.showImg1MatchedPoints()
        imageVectorExactor.extract_vectors()
        speed_matrix_list.append(imageVectorExactor.speed_matrix)

    speed_matrix_list.reverse()

    matrix = np.zeros(speed_matrix_list[0].shape, np.float32)

    for y in range(matrix.shape[0]):
        for x in range(matrix.shape[1]):
            # value = speed_matrix_list[0][y][x]

            for vector in speed_matrix_list:
                if vector[y][x][0] != 0:
                    speed_x = vector[y][x][0]
                    if speed_x > 15:
                        print(speed_x)
                    matrix[y][x][0] = vector[y][x][0]
                    matrix[y][x][1] = vector[y][x][1]
                    break

    result = np.where(matrix > 0)
    print(len(result[0]))
    print(matrix.max())

    np.save('matrix.npy', matrix)

    speed_x = np.array(speed_x)
    print(speed_x.mean())

    data = matrix[:, :, 0:1]
    data = data.reshape(700, 700)

    plt.imshow(data)
    plt.colorbar()
    plt.show()


if __name__ == '__main__':
    # show_hot()


    imageVectorExactor = ImageVectorExtractor(config.BLOCK + '%02d.png' % (0,),
                                              config.BLOCK + '%02d.png' % (1,))
    imageVectorExactor.match()
    imageVectorExactor.extract_vectors()
    imageVectorExactor.calc_not_match_vectors()
    # imageVectorExactor.showImg1MatchedPoints()
